import pandas as pd
from statsmodels.tsa.arima.model import ARIMA
from datetime import timedelta

# 1. 读取数据
fujian1 = pd.read_excel('fujian1.xlsx')

# 2. 数据预处理
# 将日期列转换为日期类型
fujian1['date'] = pd.to_datetime(fujian1['date'])

# 3. 分组和预测
result = []

# 对每个商家、商品、仓库的组合进行预测
grouped = fujian1.groupby(['seller_no', 'product_no', 'warehouse_no'])

for (seller, product, warehouse), group in grouped:
    # 取出数量数据，并按日期索引
    group.set_index('date', inplace=True)
    group = group.resample('D').sum().fillna(0)  # 以天为单位重采样，并填充缺失值

    # 构建并训练ARIMA模型
    model = ARIMA(group['qty'], order=(5, 1, 0))  # 根据需要调整参数
    model_fit = model.fit()

    # 进行预测
    forecast_index = pd.date_range(start='2023-05-15', end='2023-05-30', freq='D')
    forecast = model_fit.forecast(steps=len(forecast_index))

    # 收集结果
    for date, qty in zip(forecast_index, forecast):
        result.append({
            'seller_no': seller,
            'product_no': product,
            'warehouse_no': warehouse,
            'date': date,
            'forecast_qty': qty
        })

# 4. 将结果转换为 DataFrame
result_df = pd.DataFrame(result)

# 你可以选择将结果保存到结果文件中
# result_df.to_excel('result1.xlsx', index=False)

print(result_df)
